import cv2
import numpy as np

# 过滤小物体
min_w = 90
min_h = 90

# 检测线的高度
line_height = 550

# 线的偏移量
offset = 7

# 统计车的数量
carno = 0

# 统计车辆数组
cars = []

def center(x, y, w, h):
    x1 = int(w / 2)
    y1 = int(h / 2)
    cx = x + x1
    cy = y + y1
    return cx,cy

cap = cv2.VideoCapture('img/video.mp4')

# 祛除视频中的背景
bg_subtractor = cv2.bgsegm.createBackgroundSubtractorMOG()

# 形态学处理 kernel
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))

while True:
    ret, frame = cap.read()
    if (ret == True):
        # 转换为灰度图像
        cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        # 去噪 (高斯)
        blur = cv2.GaussianBlur(frame, (3, 3), 5)
        # 应用背景减法算法
        mask = bg_subtractor.apply(blur)
        # 腐蚀 去掉图中的小斑块
        erode = cv2.erode(mask, kernel)
        # 膨胀 还原放大
        dilate = cv2.dilate(erode, kernel, iterations=2)
        # 闭运算 去掉物体内部的小块
        close = cv2.morphologyEx(dilate, cv2.MORPH_CLOSE, kernel)
        close = cv2.morphologyEx(dilate, cv2.MORPH_CLOSE, kernel)
        # 查找轮廓
        cnts,h = cv2.findContours(close, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        # 画线
        cv2.line(frame, (0, line_height), (1280, line_height), (255, 255, 0), 3)
        
        for c in cnts:
            x, y, w, h = cv2.boundingRect(c)
            # 对车辆的宽高进行判断、验证是否有效的车辆
            isValid = (w >= min_w) and (h >= min_h)
            if not isValid:
                continue
            # 到这里都是有效的车辆
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
            cpoint = center(x, y, w, h)
            cars.append(cpoint)
            cv2.circle(frame, (cpoint), 5, (0, 0, 255), -1)
            # 统计车
            for (x,y) in cars:
                if ( ( y > line_height - offset) and (y < line_height + offset) ):
                    print('car is detected')
                    carno += 1
                    cars.remove((x,y))
        # 显示车辆数量
        cv2.putText(frame, 'car no: ' + str(carno), (500, 60), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 2)
        cv2.imshow('video', frame)
    
    key = cv2.waitKey(40)
    if key == 27:
       break

# cap 资源释放
cap.release()
# 窗口资源释放
cv2.destroyAllWindows()